#include "common.hpp"
#include "byte_model.hpp"
#include "bit_model.hpp"
#include "range_coder.hpp"
#include "qsmodel.hpp"

/*

Ideas:

Compression without much speed impact:
* Model match length within the offset word in some way.
* Larger context for literal encoding.

Compression:
+ Lazy matching.
- Two or more hashes to yield more offsets.
* Try a more accurate bit model. BitModelFast2 does not yield significantly better results while adding almost 0.2s to runtime on enwik8.
* Don't use a hash for selecting slot deterministically, search all slots for matches.

Speed:
+ Store a different hash in the offset slot and match against it first to reject matches without more memory accesses.
  This would work great with two hashtables with different offsets, as we can store hash 1 in the hash 2 slot and vice
  versa.
* Instead of encoding the slot index within the range coding stream, arrange to have it delay-written as a normal byte
  so that the decoder can read it with a simple getbyte(). This could also apply to code literals in contexts
  that are new.
* Faster memory compare in match checker.
*/

struct Rolz
{

	#define HBITS 9
	#define REJBITS 8 // Always 8 because it's fit in a byte
	#define OFFBITS 3
	#define OFFCOUNT (1<<OFFBITS)
	#define OFFCOUNTSH 1
	#define HCOUNT (1<<HBITS)


	struct Slot
	{
		Slot()
#if 0
		: offsets(0)
#else
		: next(0)
#endif
		{
			for(int i = 0; i < OFFCOUNT; ++i)
			{
				offsets[i] = 0;
				hashes[i] = 0;
			}
		}

		uint8 hashes[OFFCOUNT];
		uint32 offsets[OFFCOUNT];
		uint32 next;
	};

	Slot tab[256][HCOUNT];
	//ByteModelBitwise literalModel[256];
	//ByteModelMultiSymbol literalModel[256];
	QSModel<256> literalModel[256];
	QSModel<OFFCOUNT> offsetModel[256];
	//QSModel<1<<OFFCOUNTSH> offsetRemModel;
	QSModel<HCOUNT> hashModel[256];
	QSModel<256> lengthModel[256];
	BitModelFast bitModels[256][6];
	//BitModelApproxCounter bitModels[4*8];

	Rolz()
	{
		BitModelApproxCounter::init();
	}

	uint32 checkMatch(uint32 offset, uint8 const* begin, uint8 const* cur, uint8 const* end, uint32 minlen = 2)
	{
	/*
		if(h2 != (offset & 31))
			return 0;
			*/
		uint8 const* src = begin + offset;

		// If the offset is too close, we reject it because the decoder won't know it until later.
		// 3 would actually do, but 4 allow us to copy 4 byte at a time in the decoder without
		// risk of corrupting data.
		//if(cur - src < 4)
		//	return 0;

		// Check byte after the minimum length, it must match
		// for us to have any hope of succeeding.
		if(end - cur <= minlen
		//|| *reinterpret_cast<uint16 const*>(src+minlen-1) != *reinterpret_cast<uint16 const*>(cur+minlen-1))
		|| src[minlen] != cur[minlen])
			return 0;

		//uint32 maxlen = std::min(std::min(uint32(cur - src - 4), uint32(end - cur)), uint32(255u+3u));

		// TODO: Have a real circular 2^24 (or smaller) buffer
		// 2^24 window
		//if(cur - src >= 0x1000000)
		//	return 0; // Too far match

	#if 1
		
		uint32 len = 0;
		//for(; len < maxlen; ++len)
		for(; cur != end && *cur++ == *src++ && len < 255+3; ++len)
			/* Nothing */;

	#elif 0
		while(true)
		{
			if(cur == end || *cur++ != *src++) break;
			++len;
			if(cur == end || *cur++ != *src++) break;
			++len;
			if(cur == end || *cur++ != *src++) break;
			++len;
			if(cur == end || *cur++ != *src++) break;
			++len;
		}
	#else
		while(cur < end)
		{
			if(*reinterpret_cast<uint32 const*>(cur) != *reinterpret_cast<uint32 const*>(src))
				break;

			cur += 4;
			src += 4;
			len += 4;
		}

		while(cur != end)
		{
			if(*cur++ != *src++)
				break;

			++len;
		}
	#endif

		return len;
	}

	uint32 decodeLength(Decoder1& e)
	{
	#if 0 // TEMP
		uint32 div = 0;
		while(!bitModels[4].decode(e))
			++div;
		uint32 len = (div << 1) | e.decode_uni();
		return len + 3;
	#endif
		return 0;
	}

	// TODO: Use context for length
	void encodeLength(Coder1& e, uint32 len, QSModel<256>* lm)
	{
		len -= 3; // Minimum length of 3

	#if 0
		for(int i = 0; i < bits; ++i)
			bitModels[3].encode(e, 0);
		bitModels[3].encode(e, 1);
		e.encode_unibits(len ^ (1<<bits), bits);
	#elif 0
		// Using a uniform model here only sacrifies about 0.002 bpc on enwik8, but
		// speeds it up around 0.075s.
		for(int i = 0; i < bits; ++i)
			e.encode_uni(0);
		e.encode_uni(1);
		e.encode_unibits(len ^ (1<<bits), bits);
	#elif 0
		uint32 div = len >> 1;
		for(uint32 i = 0; i < div; ++i)
			e.encode_uni(0);
		e.encode_uni(1);
		e.encode_uni(len & 1);
	#elif 0
		uint32 div = len >> 1;
		for(uint32 i = 0; i < div; ++i)
			bmb[4].encode(e, 0);
		bmb[4].encode(e, 1);
		e.encode_uni(len & 1);
	#else
		lm->encode(e, len);
	#endif
		
	}

	#define DECODER_UPDATETABLE() { \
		if(skipUpdate <= 0) { \
			uint32 h = HASH(cur); \
			uint32 h2 = HASH2(cur); \
			Slot& slot = tab[cur & 0xff][h]; \
			char next = (slot.next - 1) & (OFFCOUNT-1); \
			slot.offsets[next] = uint32(p - begin - 3); \
			slot.hashes[next] = h2; \
			slot.next = next; \
		} else { \
			--skipUpdate; \
		} \
	}

	#define CHECKP() { if(end - p < 3) goto finish; }

	// Hash includes the last byte, but that doesn't matter
	// because it's always the same for a given context.
	//#define HASH(c) (((c & 0x00ffffff) * 2654435761) >> (32-HBITS))
	//#define HASH(c) (uint32((c) * 2654435761) >> (32-HBITS))
	//#define HASH2(c) (uint32((c) * 0x12345678) >> (32-REJBITS))
	#define HASH(c) (((c) * 3625445617) >> (32-HBITS))
	#define HASH2(c) (((c) * 2654435761) >> (32-REJBITS))

	#define READ() { \
		if(p == end) goto finish; \
		cur = (cur >> 8) | (p[3] << 24); \
		++p; \
	}

	#define PEEK3(pos) (*reinterpret_cast<uint32 const*>(pos) << 8)
	#define PEEK4(pos) (*reinterpret_cast<uint32 const*>(pos-1))

	#define WRITE(v) { \
		uint8 _v = (v); \
		*p++ = _v; \
		if(p == end) goto end; \
		cur = (cur >> 8) | ((_v) << 24); \
	}
	
	void decompress(ByteReader in, uint8* begin, uint8* end, uint8* cmp)
	{
		Decoder1 e(in);

		uint8* p = begin;

		uint32 cur = 0;
		uint32 last = 0;
		int skipUpdate = 2; // First two should not cause table updates as they would be invalid offsets

		while(true)
		{
			uint8 last = cur >> 24;

			BitModelFast* bmb = bitModels[last];

			if(!bmb[0].decode(e))
			{
				// Match
				int offset = bmb[1].decode(e);
				offset |= bmb[2].decode(e) << 1;

				int h = e.decode_unibits(HBITS);
				uint32 len = decodeLength(e);

				//offset = (offset + tab[last][h].next) & (OFFCOUNT-1);
				
				Slot& slot = tab[last][h];

				uint32 offs = (slot.offsets[offset] >> REJBITS);
				uint8 const* src = begin + offs;

				// TODO: Safety checks here

				assert(offs < p - begin - 3);
				assert(offs + len < end - begin);

				assert(p + len <= end);

				// Update all tables
				WRITE(*src++);
				DECODER_UPDATETABLE(); // -2 before match
				WRITE(*src++);
				DECODER_UPDATETABLE(); // -1 before match
				WRITE(*src++);
				DECODER_UPDATETABLE(); // At match

				// Copy bytes [3, len)

				len -= 3;

				//if(len != 0)
				{
				
					while(len >= 4)
					{
						*reinterpret_cast<uint32*>(p) = *reinterpret_cast<uint32 const*>(src);
						p += 4;
						src += 4;
						len -= 4;
					}

					while(len-- > 0)
						*p++ = *src++;

					//p += len - 3;
					cur = PEEK4(p - 3);
				}
				if(p == end) goto end;
				skipUpdate = 2; // Encoder does not update at -2, -1

			}
			else
			{
				int lit = literalModel[last].decode(e);

				WRITE(lit);
				DECODER_UPDATETABLE();
			}
		}
		
	end:
		return;
	}

	#if 0
	#define ENCODER_UPDATETABLE(pos) { \
		char next = (slot.next - 1) & (OFFCOUNT-1); \
		slot.offsets[next] = (uint32((pos) - begin) << REJBITS) | h2; \
		slot.next = next; \
	}
	#else
	#define ENCODER_UPDATETABLE(pos) do { \
		uint8 const* _pos = (pos); \
		uint32 dist = _pos - begin; \
		if(dist == prevupdate) printf("Duplicate update @ %d\n", dist); \
		prevupdate = dist; \
		uint32 cur2 = PEEK4(_pos); \
		uint32 h_ = HASH(cur2); \
		uint32 h2_ = HASH2(cur2); \
		Slot& slot = tab[cur2 & 0xff][h_]; \
		int idx = --slot.next & (OFFCOUNT-1); \
		slot.offsets[idx] = uint32(_pos - begin); \
		slot.hashes[idx] = h2_; \
	} while(0)
	#endif

	#define LAZYPARSE 1

	void compress(uint8 const* begin, uint8 const* end, ByteWriter out, std::vector<uint8>& testout)
	{
		Coder1 e(out);

		uint8 const* p = begin;
		uint8 last = 0;

		uint32 lazyMatchLen = 0;
		int lazyOffset = 0;
		uint8 lazyLast = 0;
		uint32 lazyH = 0;
		int lit = 0;

		int prevupdate = -1;

		CHECKP();
		uint32 cur;
		cur = PEEK3(p);

		for(int i = 0; i < 4; ++i)
		{
			last = cur & 0xff;
			bitModels[last][0].encode(e, 1);
			literalModel[last].encode(e, (cur >> 8) & 0xff);
			READ();
		}

		while(true)
		{
			uint32 h, h2, last;
			h = HASH(cur);
			h2 = HASH2(cur);
			last = cur & 0xff;

			Slot& slot = tab[last][h];
			
			uint32 max = 0;
			//uint32 min = lazyMatchLen + 2;
			uint32 maxi = 0;

			int rot = slot.next;

			for(int i = 0; i < OFFCOUNT; ++i)
			{
				int idx = (i + rot) & (OFFCOUNT-1);
				if(slot.hashes[idx] != h2)
					continue;

				uint32 offset = slot.offsets[idx];
				uint32 m = checkMatch(offset, begin, p, end, max);
				if(m > max)
				{
					maxi = i;
					max = m;
					//min = m;
				}
			}

			ENCODER_UPDATETABLE(p - 3);

			BitModelFast* bmb = bitModels[last];

#if LAZYPARSE
			if(lazyMatchLen)
			{
			#if 1
				if(max > lazyMatchLen + 1)
				{
					bmb[0].encode(e, 1);
					literalModel[lazyLast].encode(e, cur & 0xff);
					READ();
					//testout.push_back(cur & 0xff);

					lazyMatchLen = max;
					lazyOffset = maxi;
					lazyH = h;
					lazyLast = last;
				}
				else
			#endif
				{
					ENCODER_UPDATETABLE(p - 2);
					if(lazyMatchLen > 3)
						ENCODER_UPDATETABLE(p - 1);
					
					int offs = lazyOffset;
					bmb[0].encode(e, 0);
				#if 0
					bitModels[1].encode(e, offs & 1);
					bitModels[2].encode(e, offs >> 1);
				#elif 1
					offsetModel[last].encode(e, offs);
				#else
					{
						int div = offs >> OFFCOUNTSH;
						int const max = ((OFFCOUNT-1) >> OFFCOUNTSH);
						for(int i = 0; i < div; ++i)
							e.encode_uni(0);
						if(div < max)
							e.encode_uni(1);
						e.encode_unibits(offs & ((1<<OFFCOUNTSH)-1), OFFCOUNTSH);
						//offsetRemModel.encode(e, offs & ((1<<OFFCOUNTSH)-1));
					}
				#endif

				#if 0
					e.encode_unibits(lazyH, HBITS);
				#else
					hashModel[last].encode(e, lazyH);
				#endif
					encodeLength(e, lazyMatchLen, lengthModel + last);

					p += lazyMatchLen - 1; // 1 byte lookahead

					CHECKP();
					cur = PEEK4(p);

					lazyMatchLen = 0;
				}
			}
			else
#endif
			if(max >= 3)
			{
#if LAZYPARSE
				// Update here because it's the start of a match or a literal
				//ENCODER_UPDATETABLE(p);
				lazyMatchLen = max;
				lazyOffset = maxi;
				lazyH = h;
				lazyLast = last;
				READ();
#else
				ENCODER_UPDATETABLE(p - 2);
				ENCODER_UPDATETABLE(p - 1);
				if(max > 3)
					ENCODER_UPDATETABLE(p);

				QSModel<256>* lmb = &lengthModel[last];

				bmb[0].encode(e, 0);
				offsetModel[last].encode(e, maxi);
				hashModel[last].encode(e, h);
				//e.encode_unibits(h, HBITS);
				encodeLength(e, max, lmb);

				p += max;

				CHECKP();
				cur = PEEK4(p);
#endif
			}
			else
			{
				//ENCODER_UPDATETABLE(p);
				bmb[0].encode(e, 1);
				// Encoding literals as unibits makes the time on enwik8, 2.7s -> 2s or so.
				// We should see if it's possible to optimize the literal encoding (mostly the model).
				// Using an order-0 model cost about 0.1 bpc on enwik8.
				literalModel[last].encode(e, (cur >> 8) & 0xff);
				//e.encode_unibits((cur >> 8) & 0xff, 8);
				//testout.push_back(cur >> 8);
			do_read:
				READ();
			}
		}

	finish:
		// There may be up to 3 bytes left here

		while(p != end)
		{
			bitModels[last][0].encode(e, 1);
			literalModel[last].encode(e, *p);
			last = *p++;
		}

		e.flush();
	}
};

void rolz_compress(uint8 const* begin, uint8 const* end, ByteWriter out, std::vector<uint8>& testout)
{
	std::auto_ptr<Rolz> r(new Rolz);

	r->compress(begin, end, out, testout);
}

void rolz_decompress(ByteReader in, uint8* begin, uint8* end, uint8* cmp)
{
	std::auto_ptr<Rolz> r(new Rolz);

	r->decompress(in, begin, end, cmp);
}